#! /usr/bin/env python
#coding=utf-8

import os
import sys
import concurrent.futures
import threading
import sqlite3

from tablebuilder import MemoryRecordBuilder, ProcessTableBuilder, SystemObjectTableBuilder, RegionTypeTableBuilder, RegionTableBuilder

sys.path.append(os.path.join(os.path.dirname(os.path.realpath(__file__)), "../"))
from vmmap import ProcessMgr

class MemDbBuilder(ProcessMgr):
	def __init__(self, db_file, mem_dir, description=""):
		ProcessMgr.__init__(self, mem_dir)
		self.__db_file = db_file
		self.__conn, self.__cursor = self.connectDB()

		name = os.path.basename(mem_dir)
		self.__add_record(name, description)

	def get_record_id(self):
		return self.__id

	def __add_record(self, name, description):
		builder = MemoryRecordBuilder(self.__cursor)
		builder.createTable(False)

		info = self.get_meminfo()
		info["name"] = name
		info["description"] = description
		builder.saveToDB(info, no_id=True)
		self.__id = self.__cursor.lastrowid

	def __process_preprocessor(proc, idx, mgr):
		proc["record_id"] = mgr.get_record_id()

	def __region_type_preprocessor(obj, idx, mgr):
		obj["record_id"] = mgr.get_record_id()
		obj["process_id"] = obj["__process"]["id"]

	def __region_preprocessor(obj, idx, mgr):
		obj["record_id"] = mgr.get_record_id()
		proc = obj["__process"]
		obj["process_id"] = proc["id"]
		obj["type_id"] = proc["__region_types_dict"][obj["type"]]

	def __system_object_preprocessor(obj, idx, mgr):
		obj["record_id"] = mgr.get_record_id()

	def __threadpool_initializer(local, mgr, builderClass):
		local.conn = sqlite3.connect(mgr.__db_file, timeout=50)
		local.cursor = local.conn.cursor()
		local.mgr = mgr
		local.builder = builderClass(local.cursor)

	def threadpool_do_tasks(self, task, items, builderClass):
		local = threading.local()
		executor = concurrent.futures.ThreadPoolExecutor(initializer=MemDbBuilder.__threadpool_initializer,
														 initargs=(local, self, builderClass))
		futures = [executor.submit(task, local, item) for item in items]
		failCnt = 0
		for future in concurrent.futures.as_completed(futures):
			result = future.result()
			if not result:
				failCnt = failCnt + 1
		print("threadpool failed %d for %d jobs" % (failCnt, len(items)))
		executor.shutdown()

	def __load_process_ids_task(local, proc):
		name = proc["command"].replace("'", "")
		local.cursor.execute("select id from %s where command='%s'" % (local.builder.getTableName(), name))
		for row in local.cursor:
			proc["id"] = row[0]
			break
		return True

	def __load_system_object_ids_task(local, obj):
		name = obj["name"].replace("'", "")
		local.cursor.execute("select id from %s where name='%s'" % (local.builder.getTableName(), name))
		for row in local.cursor:
			obj["id"] = row[0]
			break

	def addProcesses(self):
		print("Add processes now ...")
		builder = ProcessTableBuilder(self.__cursor)
		if builder.createTable(False):
			self.__cursor.execute("DROP INDEX IF EXISTS index_process_names")
			self.__cursor.execute("CREATE INDEX index_process_names ON processes (command)")

			self.__cursor.execute("drop view if exists processes_details")
			cols = builder.getColumns()
			cols = ["processes.%s as %s" % (k, k) for k in cols]
			sqlcmd = "create view IF NOT EXISTS [processes_details] as SELECT mem_records.name as record_name, processes.id as id, %s from processes LEFT OUTER join mem_records on processes.record_id=mem_records.id" % ",".join(cols)
			self.__cursor.execute(sqlcmd)
		builder.addMultiObjToDb(self.get_all(), self, MemDbBuilder.__process_preprocessor, no_id=True)
		self.threadpool_do_tasks(MemDbBuilder.__load_process_ids_task, self.get_all(), ProcessTableBuilder)

	def __add_region_types_task(local, proc):
		local.builder.addMultiObjToDb(proc["__region_types"], local.mgr, MemDbBuilder.__region_type_preprocessor, no_id=True)
		return True

	def addRegionTypes(self):
		print("Add region types now ...")
		builder = RegionTypeTableBuilder(self.__cursor)
		if builder.createTable(False):
			self.__cursor.execute("DROP INDEX IF EXISTS index_region_types_names")
			self.__cursor.execute("CREATE INDEX index_region_types_names ON region_types (name)")

			self.__cursor.execute("drop view if exists region_types_details")
			cols = builder.getColumns()
			cols = ["region_types.%s as %s" % (k, k) for k in cols]
			self.__cursor.execute("CREATE VIEW [region_types_details] as SELECT mem_records.name as record_name, processes.command as command, region_types.id as id,%s from region_types LEFT OUTER join mem_records on region_types.record_id=mem_records.id LEFT OUTER join processes on region_types.process_id=processes.id" % ",".join(cols))
		self.threadpool_do_tasks(MemDbBuilder.__add_region_types_task, self.get_all(), RegionTypeTableBuilder)

		print("Load region types for each process now ...")
		for proc in self.get_all():
			sqlcmd = "select id, name from region_types where process_id=%d" % proc["id"]
			self.__cursor.execute(sqlcmd)
			proc["__region_types_dict"] = {}
			for row in self.__cursor:
				proc["__region_types_dict"][row[1]] = row[0]

	def __add_regions_task(local, proc):
		print("Add %d regions for proc %d" % (len(proc["__regions"]), proc["pid"]))
		local.builder.addMultiObjToDb(proc["__regions"], local.mgr, MemDbBuilder.__region_preprocessor, no_id=True)
		return True

	def addRegions(self):
		print("Add regions now ...")
		builder = RegionTableBuilder(self.__cursor)
		if builder.createTable(False):
			self.__cursor.execute("DROP INDEX IF EXISTS index_region_details")
			self.__cursor.execute("CREATE INDEX index_region_details ON regions (detail)")

			self.__cursor.execute("drop view if exists regions_details")
			cols = builder.getColumns()
			cols = ["regions.%s as %s" % (k, k) for k in cols]
			sqlcmd = "create view IF NOT EXISTS [regions_details] as SELECT mem_records.name as record_name, processes.command as command, region_types.name as type, regions.id as id, %s from regions LEFT OUTER join mem_records on regions.record_id=mem_records.id LEFT OUTER join processes on regions.process_id=processes.id LEFT OUTER join region_types on regions.type_id=region_types.id" % ",".join(cols)
			self.__cursor.execute(sqlcmd)

		self.threadpool_do_tasks(MemDbBuilder.__add_regions_task, self.get_all(), RegionTableBuilder)

	def addSystemObjects(self):
		builder = SystemObjectTableBuilder(self.__cursor)
		builder.createTable()
		builder.addMultiObjToDb(self.get_all_objects(), self, MemDbBuilder.__system_object_preprocessor, no_id=True)
		self.threadpool_do_tasks(MemDbBuilder.__add_regions_task, self.get_all(), RegionTableBuilder)

	def addProcessDetails(self):
		builder = SystemObjectTableBuilder(self.__cursor)
		builder.createTable()
		builder.addMultiObjToDb(self.get_all_objects(), no_id=True)

	def addAllMemInfo(self):
		self.addProcesses()
		self.addRegionTypes()
		self.addRegions()

	def connectDB(self):
		# must set a big timeout value, otherwise concurrent db operation may fail
		conn = sqlite3.connect(self.__db_file, timeout=20)
		cursor = conn.cursor()
		return conn, cursor

	def close_db(self):
		self.__conn.commit()
		self.__cursor.close()
		self.__conn.close()

if __name__ == "__main__":
	builder = MemDbBuilder("symdb.db", "/Users/handy/Documents/work/projects/python/macArchInfo/collector/vmmap/backlog_20230513_2340")

	builder.addAllMemInfo()

	builder.close_db()
